import json
import os
import pickle

import numpy

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

import operator

from torch.utils.data import Dataset
import pandas as pd
import torch
import torch.nn as nn
import numpy as np
from vedo import *
from scipy.spatial import distance_matrix

from utils import color_2_label
from meshsegnet import *

class Mesh_Dataset(Dataset):
    def __init__(self, data_list_path, patch_size=6000, num_seg_classes=15):
        self.data_list = pd.read_csv(data_list_path, header=None)
        self.patch_size = patch_size
        self.num_seg_classes = num_seg_classes
        
        with open('config.json', 'r') as f:
            config = json.load(f)
            print("load_pkl")
            
        mod = config["mod"]
        dataset_path = config["data"]["dataset_path"]
        seg_matrix_path = config["data"]["seg_matrix_path"] + '{}.pkl'.format(mod)
        seg_matrix_path = dataset_path + seg_matrix_path
        
        
        with open(seg_matrix_path, 'rb') as fr:
            seg_matrix_map = pickle.load(fr)
        self.seg_matrix_map = seg_matrix_map

    def __len__(self):
        return self.data_list.shape[0]

    def __getitem__(self, idx):
        with open('config.json', 'r') as f:
            config = json.load(f)

        mod = config["mod"]
        dataset_path = config["data"]["dataset_path"]
        label_path = dataset_path + config["data"]["label_data_path"] + mod + ".json"

        with open(label_path, 'r') as f:
            labels_cls_map = json.load(f)

        if torch.is_tensor(idx):
            idx = idx.tolist()

        i_mesh = self.data_list.iloc[idx][0]

        mesh = Mesh(i_mesh)

        colors = mesh.cellcolors[:, 0:3] / 255.
        points = mesh.points()
        mean_cell_centers = mesh.center_of_mass()
        points[:, 0:3] -= mean_cell_centers[0:3]

        ids = np.array(mesh.faces())
        cells = points[ids].reshape(mesh.ncells, 9).astype(dtype='float32')

        mesh.compute_normals()
        normals = mesh.celldata['Normals']

        barycenters = mesh.cell_centers()
        barycenters -= mean_cell_centers[0:3]

        maxs = points.max(axis=0)
        mins = points.min(axis=0)
        means = points.mean(axis=0)
        stds = points.std(axis=0)
        nmeans = normals.mean(axis=0)
        nstds = normals.std(axis=0)

        for i in range(3):
            cells[:, i] = (cells[:, i] - means[i]) / stds[i]
            cells[:, i + 3] = (cells[:, i + 3] - means[i]) / stds[i]
            cells[:, i + 6] = (cells[:, i + 6] - means[i]) / stds[i]
            barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
            normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]

        X = np.column_stack((cells, barycenters, normals, colors))

        # A_S = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
        # A_L = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
        # D = distance_matrix(X[:, 9:12], X[:, 9:12])
        # A_S[D < 0.1] = 1.0
        # A_S = A_S / np.dot(np.sum(A_S, axis=1, keepdims=True), np.ones((1, X.shape[0])))
        #
        # A_L[D < 0.2] = 1.0
        # A_L = A_L / np.dot(np.sum(A_L, axis=1, keepdims=True), np.ones((1, X.shape[0])))

        # seg_matrix_path = config["data"]["seg_matrix_path"] + '{}.pkl'.format(mod)
        # seg_matrix_path = dataset_path + seg_matrix_path

        # with open(seg_matrix_path, 'rb') as fr:
        #     seg_matrix_map = pickle.load(fr)

        # print("ready: ", i_mesh)
        if i_mesh in self.seg_matrix_map.keys():
            seg_matrix = self.seg_matrix_map[i_mesh]
            # print("load: ", i_mesh)
        else:
            models_seg_path = config["data"]["models_seg_path"]
            models_seg_name = config["data"]["models_seg_name"] + '{}.tar'.format(mod)

            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            model_seg = MeshSegNet(num_classes=15, num_channels=15).to(device, dtype=torch.float)

            checkpoint = torch.load(os.path.join(models_seg_path, models_seg_name), map_location='cpu')
            model_seg.load_state_dict(checkpoint['model_state_dict'])

            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.enabled = True

            model_seg.eval()

            X_seg = torch.from_numpy(X[:, 0:15]).to(device, dtype=torch.float)
            X_seg = X_seg.transpose(1, 0)
            X_seg = X_seg.reshape([-1, X_seg.shape[0], X_seg.shape[1]])
            A_S = A_S.reshape([1, A_S.shape[0], A_S.shape[1]])
            A_L = A_L.reshape([1, A_L.shape[0], A_L.shape[1]])
            A_S = torch.from_numpy(A_S).to(device, dtype=torch.float)
            A_L = torch.from_numpy(A_L).to(device, dtype=torch.float)

            # print(X_seg.shape, A_S.shape, A_L.shape)

            # seg_matrix = None

            with torch.no_grad():
                seg_output = model_seg(X_seg, A_S, A_L).to(device, dtype=torch.float)
                seg_matrix = seg_output.cpu().numpy()
            seg_matrix_map[i_mesh] = seg_matrix.tolist()

            with open(seg_matrix_path, 'wb') as fw:
                pickle.dump(seg_matrix_map, fw)
                print('dumped: {}'.format(i_mesh))

        labels = labels_cls_map[i_mesh.split('_')[-1].split('.')[0]]
        labels = np.asarray(labels).astype(dtype='int64')

        Y = labels

        # print(Y.shape)

        X_train = np.zeros([self.patch_size, X.shape[1]], dtype='float32')
        Y_train = Y
        # Y_train = np.zeros([self.patch_size, Y.shape[1]], dtype='int32')

        valid_idx = np.array([i for i in range(X.shape[0])])
        selected_idx = np.sort(np.random.choice(valid_idx, size=self.patch_size, replace=False))

        X_train[:] = X[selected_idx, :]
        # Y_train[:] = Y[selected_idx, :]
        # seg_matrix = seg_matrix.reshape([seg_matrix.shape[1], seg_matrix.shape[2]])
        seg_matrix = seg_matrix[selected_idx, :]

        X_train = X_train.transpose(1, 0)
        # Y_train = Y_train.transpose(1, 0)

        sample = {
            'cells': torch.from_numpy(X_train),
            'labels': torch.from_numpy(Y_train),
            'seg_matrix': torch.from_numpy(seg_matrix)
        }

        return sample


if __name__ == '__main__':
    dataset = Mesh_Dataset('./list_upper/train_list_1.csv')
    sample = dataset.__getitem__(0)
    cells = sample["cells"]
    labels = sample["labels"]
    seg_matrix = sample["seg_matrix"]

    # with open('test.txt', 'w') as f:
    #     f.write(str([round(i, 4) for i in cells.tolist()]) + '\n')
    #     f.write(str([round(i, 4) for i in labels.tolist()]) + '\n')
    #     f.write(str([round(i, 4) for i in seg_matrix.tolist()]) + '\n')

    # print(cells.shape)
    # print(cells)
    # print(labels.shape)
    # print(labels)
    print(seg_matrix.shape)
    print(seg_matrix)